#!/usr/bin/python3
import os
import argparse
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(
        description='Face Norm')
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/FaceRecord/FaceBaiduBoxFeature.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data1/data/FaceRecord/FaceBaiduFeature.dat", help="path to feature file", type=str
    )
    parser.add_argument('--head_offset', type=int, default=8, help='feature head offset')
    parser.add_argument('--skip_head', action='store_true', help='skip head')
    parser.add_argument('--mode', type=int, default=0, help='mode')
    parser.add_argument('--feat_dim', type=int, default=256, help='feature dim, when mode != 0')
    args = parser.parse_args()
    return args


args = parse_args()
feat_head_offset = args.head_offset
face_feature_list_file = args.file_list
face_feature_data_file = args.feature_file


def main():
    if args.mode == 0:
        with open(face_feature_list_file, 'r') as file:
            file_count = 0
            save_file_count = 0
            file_infos = file.readlines()
            total_file_count = len(file_infos)
            info = file_infos[0].strip().split()
            feat_size = int(info[1].strip().split(',')[-1])
            feat_dim = (feat_size - feat_head_offset) // 4
            feat_file = open(face_feature_data_file, 'rb')
            out_feat_file = open(os.path.splitext(face_feature_data_file)[0]+'.norm.dat', 'wb')
            out_file = open(os.path.splitext(face_feature_list_file)[0]+'.norm.txt', 'w')
            for line in file_infos:
                file_count += 1
                if file_count % 1000 == 0:
                    print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))

                feat_data = feat_file.read(feat_size)
                no_head_feat_data = feat_data[feat_head_offset:]
                feat = np.frombuffer(no_head_feat_data, dtype=np.float32, count=feat_dim)
                feat = np.asarray(feat) / np.linalg.norm(feat, axis=0, keepdims=True)
                if args.skip_head:
                    feat.tofile(out_feat_file)
                    idx = line.rfind(',')
                    line = line[:idx] + ',{}\n'.format(feat_dim*4)
                    out_file.write(line)
                else:
                    out_feat_file.write(feat_data[:feat_head_offset])
                    out_feat_file.write(feat.tobytes())
                    out_file.write(line)
                save_file_count += 1
            out_file.close()
            out_feat_file.close()
            feat_file.close()
            print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))
    else:
        with open(face_feature_list_file, 'r') as file:
            file_count = 0
            save_file_count = 0
            file_infos = file.readlines()
            total_file_count = len(file_infos)
            feat_dim = args.feat_dim
            feat_file = open(face_feature_data_file, 'rb')
            out_feat_file_path = os.path.splitext(face_feature_data_file)[0] + '.norm.dat'
            out_feat_file = open(out_feat_file_path, 'wb')
            for line in file_infos:
                file_count += 1
                if file_count % 1000 == 0:
                    print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))

                feat = np.fromfile(feat_file, dtype=np.float32, count=feat_dim)
                feat = np.asarray(feat) / np.linalg.norm(feat, axis=0, keepdims=True)
                feat.tofile(out_feat_file)
                save_file_count += 1
            out_feat_file.close()
            feat_file.close()
            os.system('chmod a+wr {}'.format(out_feat_file_path))
            print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))


if __name__ == '__main__':
    main()
